import torch
import torchvision.models as models
import time
from torchvision.transforms import ToTensor
from PIL import Image

image = Image.new('RGB', (224, 224),color=(255,0,0))
transform = ToTensor()
input_image = transform(image).unsqueeze(0)

# 确保模型和输入都在CPU上
model_static = models.resnet18(pretrained=True)
model_static.eval()

# 配置量化参数（FBGEMM是x86 CPU的量化后端）
model_static.qconfig = torch.quantization.get_default_qconfig("fbgemm")

# 准备量化 - 注意：prepare_qat用于训练时量化，如果是推理应该用prepare
model_prepared = torch.quantization.prepare(model_static, inplace=False)  # 改为prepare

# 转换为量化模型
model_quantized = torch.quantization.convert(model_prepared, inplace=False)

# 确保输入数据在CPU上
input_image = input_image.cpu()  # 如果input_image在GPU上，先移到CPU

start = time.time()

with torch.no_grad():
    out = model_quantized(input_image)  # 现在所有操作都在CPU上

print(f"推理时间: {time.time() - start}秒")